package com.cw.recommend.offline

import com.cw.recommend.common.Runner.{RunnerConfig, runSpark}
import com.cw.recommend.common.constant.{MODEL_PATH, RATING_COLLECTION}
import com.cw.recommend.common.model.UserRating
import com.cw.recommend.common.util.FileUtil
import com.cw.recommend.common.util.MongoDBUtil.readMongoDB
import com.cw.recommend.common.util.SyntaxUtil._
import org.apache.spark.mllib.recommendation.{ALS, MatrixFactorizationModel, Rating}
import org.apache.spark.rdd.RDD

import java.io.File

object ALStrainer {
  implicit val runnerConfig = RunnerConfig("als_trainer")

  def main(args: Array[String]): Unit = {
    runSpark { spark =>
      import spark.implicits._

      // 划分训练集
      val userRating = readMongoDB(spark, RATING_COLLECTION)
        .as[UserRating]
        .rdd
        .map(_.toMlLibRating)

      val splitData = userRating.randomSplit(Array(0.8, 0.2))
      val (trainSet, testSet) = (splitData(0), splitData(1))

      // 模型调参
      val ((model, rank, iter, lambda), rmse) = adjustParams(trainSet, testSet)

      if (new File(MODEL_PATH).exists())
        FileUtil.deleteDir(MODEL_PATH)
      model.save(spark.sparkContext, MODEL_PATH)

      println(s"rank: $rank, iter: $iter, lambda: $lambda => rmse: $rmse")
    }


  }


  def adjustParams(trainSet: RDD[Rating], testSet: RDD[Rating]): ((MatrixFactorizationModel, Int, Int, Double), Double) = {
    val rankSeq = 5 to 20 by 10
    val lambdaSeq = Seq(1, 0.1, 0.01)
    val iterSeq = 10 to 20 by 10
    val seq = for {
      rank <- rankSeq
      iter <- iterSeq
      lambda <- lambdaSeq
      model = ALS.train(trainSet, rank, iter, lambda)
      testData = testSet.map(r => r.user -> r.product)
    } yield {
      // 预测
      val predictSet = model.predict(testData)
      // 评估
      (model, rank, iter, lambda) -> RMSE(predictSet, testSet)
    }
    seq.sortWith { case ((_, a), (_, b)) => a > b }
      .head


  }

  private def RMSE(predictSet: RDD[Rating], testSet: RDD[Rating]): Double = {
    val keyFunc = (r: Rating) => r.user -> r.product
    predictSet
      .keyBy(keyFunc)
      .join(testSet.keyBy(keyFunc))
      .map { case (_, (a, b)) => math.pow(2, a.rating - b.rating) }
      .mean |>
      math.sqrt
  }


}
